#!/usr/bin/env python
# coding=utf-8
import os
import subprocess
import json
import pandas as pd
import tkinter as tk
from tkinter import filedialog, messagebox, Checkbutton

SUPPORTED_EXTENSIONS = ['.mp4', '.avi', '.mov', '.mkv', '.flv', '.wmv', '.ari', '.dng', '.mxf',
                        '.r3d', '.arw', '.jpg', '.jpeg', '.dpx', '.cine', '.raw']

def get_video_info_ffmpeg(file_path):
    try:
        cmd = f"ffprobe -v error -print_format json -show_format -show_streams \"{file_path}\""
        # 使用 universal_newlines 而不是 text
        result = subprocess.run(cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, universal_newlines=True)
        if result.returncode != 0:
            print(f"FFprobe Error for {file_path}: {result.stderr}")
            return None

        info_json = json.loads(result.stdout)
        video_stream = next((stream for stream in info_json['streams'] if stream['codec_type'] == 'video'), None)
        audio_streams = [stream for stream in info_json['streams'] if stream['codec_type'] == 'audio']

        frame_rate = round(
            float(video_stream['r_frame_rate'].split('/')[0]) / float(video_stream['r_frame_rate'].split('/')[1]),
            2) if video_stream else '未知'

        file_size = os.path.getsize(file_path) / (1024 * 1024)  # 文件大小MB
        file_size_decimal = os.path.getsize(file_path) / 1000000  # 文件大小十进制MB

        audio_info = audio_streams[0] if audio_streams else {}

        return {
            '文件名': os.path.basename(file_path),
            '帧率': frame_rate,
            '持续时间秒': float(video_stream['duration']) if video_stream and 'duration' in video_stream else '未知',
            '文件大小MB': file_size,
            '文件大小十进制MB': file_size_decimal,
            '音频流数量': len(audio_streams),
            '音频编码': audio_info.get('codec_name', '未知') if audio_info else '未知',
            '采样率': audio_info.get('sample_rate', '未知') if audio_info else '未知',
            '声道数': audio_info.get('channels', '未知') if audio_info else '未知',
            '文件路径': file_path
        }
    except Exception as e:
        print(f"Error getting video info using ffprobe for file {file_path}: {e}")
    return None

def process_videos_in_directory(directory_path, ignored_extensions):
    video_files = [f for f in os.listdir(directory_path) if os.path.splitext(f)[1].lower() in SUPPORTED_EXTENSIONS and os.path.splitext(f)[1].lower() not in ignored_extensions]
    video_infos = [get_video_info_ffmpeg(os.path.join(directory_path, file)) for file in video_files]
    return video_infos

def save_to_excel(data, output_path):
    df = pd.DataFrame(data)
    expected_columns = ['文件名', '帧率', '持续时间秒', '文件大小MB', '文件大小十进制MB', '音频流数量', '音频编码', '采样率', '声道数', '文件路径']
    for col in expected_columns:
        if col not in df.columns:
            df[col] = None
    df = df[expected_columns]
    df.to_excel(output_path, index=False)
def setup_gui():
    root = tk.Tk()
    root.title("路径视频信息提取脚本")
    root.geometry("700x230")  # 增加窗口高度以适应新的UI布局

    folder_path_var = tk.StringVar()
    output_path_var = tk.StringVar()
    ignored_formats_vars = {ext: tk.BooleanVar() for ext in SUPPORTED_EXTENSIONS}

    folder_label = tk.Label(root, text="选择视频文件夹:", font=("Arial", 12))
    folder_label.grid(row=0, column=0, padx=5, pady=5)
    folder_entry = tk.Entry(root, textvariable=folder_path_var, width=50, font=("Arial", 12))
    folder_entry.grid(row=0, column=1, padx=5, pady=5)
    folder_button = tk.Button(root, text="浏览", command=lambda: folder_path_var.set(filedialog.askdirectory()), height=1, width=10, font=("Arial", 12))
    folder_button.grid(row=0, column=2, padx=5, pady=5)

    output_label = tk.Label(root, text="选择Excel文件保存位置:", font=("Arial", 12))
    output_label.grid(row=1, column=0, padx=5, pady=5)
    output_entry = tk.Entry(root, textvariable=output_path_var, width=50, font=("Arial", 12))
    output_entry.grid(row=1, column=1, padx=5, pady=5)
    output_button = tk.Button(root, text="浏览", command=lambda: output_path_var.set(filedialog.asksaveasfilename(defaultextension=".xlsx", filetypes=[("Excel文件", "*.xlsx")])), height=1, width=10, font=("Arial", 12))
    output_button.grid(row=1, column=2, padx=5, pady=5)

    ignore_label = tk.Label(root, text="是否有需要忽略的封装格式:", font=("Arial", 12))
    ignore_label.grid(row=2, column=0, columnspan=3, padx=5, pady=5)

    checkboxes_frame = tk.Frame(root)
    checkboxes_frame.grid(row=3, column=0, columnspan=3, padx=5, pady=5)
    half_point = len(SUPPORTED_EXTENSIONS) // 2
    for i, (ext, var) in enumerate(ignored_formats_vars.items()):
        check = Checkbutton(checkboxes_frame, text=ext, var=var, font=("Arial", 11))
        if i < half_point:
            check.grid(row=0, column=i, sticky='w', padx=2, pady=2)
        else:
            check.grid(row=1, column=i - half_point, sticky='w', padx=2, pady=2)

    def start_processing():
        directory_path = folder_path_var.get()
        output_path = output_path_var.get()
        ignored_extensions = [ext for ext, var in ignored_formats_vars.items() if var.get()]
        if directory_path and output_path:
            video_infos = process_videos_in_directory(directory_path, ignored_extensions)
            save_to_excel(video_infos, output_path)
            messagebox.showinfo("完成", "Excel文件已生成")
        else:
            messagebox.showwarning("警告", "请确保已选择视频文件夹和输出文件路径")

    extract_button = tk.Button(root, text="提取信息", command=start_processing, height=2, width=20, font=("Arial", 12))
    extract_button.grid(row=4, column=1, padx=5, pady=10)  # 增加了按钮的行数

    root.mainloop()

if __name__ == "__main__":
    setup_gui()
